# Author: WY
# Date: 2022/11/10 11:09

from sklearn import preprocessing
from torch.utils.data import Dataset, DataLoader
import numpy as np
import matplotlib
import matplotlib.pylab as plt

# Dataset
class my_Dataset(Dataset):
    def __init__(self, data):
        self.data = data

    def __len__(self):
        return len(self.data)

    def __getitem__(self, item):
        return self.data[item]


# 数据预处理MinMaxScaler(数据范围缩放到[-1, 1]之间)
minmaxScaler = preprocessing.MinMaxScaler(feature_range=(-1., 1.))


def get_loader(split: float, batch_size: int, data):
    # 取出标签值
    # label_data = data[:, 0].reshape(-1, 1)

    # 将特征值进行归一化
    scaled_data = minmaxScaler.fit_transform(data)
    # 将标签值合并到数组
    # scaled_data = np.concatenate((scaled_data, label_data), axis=1)

    data_len = len(data)
    split_line = int(data_len * split)
    # DataSet和DataLoader
    dataset_train = my_Dataset(scaled_data[0: split_line])
    dataset_test = my_Dataset(scaled_data[split_line:])
    dataloader_train = DataLoader(dataset_train, batch_size=batch_size, shuffle=False, drop_last=True)
    dataloader_test = DataLoader(dataset_test, batch_size=batch_size, shuffle=False, drop_last=True)
    return dataloader_train, dataloader_test

